import matplotlib

matplotlib.use('TkAgg')

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import ast

df = pd.read_excel('./results.xlsx')
dict_list = df.to_dict(orient='records')

tasks = df['task'].unique()
models = df['model'].unique()

# models = ['CRISP', 'Virchow2', 'Virchow', 'UNI', 'ours']
# tasks = ['ER', 'Grade', 'HER2', 'Ki67', 'pCR', 'PR']
colors = ['#5B4E9D', '#E74C9C', '#FF6B9D', '#FF8B6B', '#FFB84D', '#F4E04D']

roc_data = {model: {} for model in models}

for task in tasks:
    task_data = df[df['task'] == task]
    for model in models:
        value_df = task_data[task_data['model'] == model]
        roc_data[model]['auc'] = value_df['auc'].values[0]
        roc_data[model]['tpr'] = np.array(ast.literal_eval(value_df['tpr'].values[0]))
        roc_data[model]['fpr'] = np.array(ast.literal_eval(value_df['fpr'].values[0]))

# 直接定义每个方法的ROC曲线数据（FPR和TPR）
# roc_data = {
#     'CRISP': {
#         'fpr': np.array([0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5,
#                          0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0]),
#         'tpr': np.array([0.0, 0.35, 0.55, 0.68, 0.76, 0.82, 0.86, 0.89, 0.91, 0.93, 0.94,
#                          0.95, 0.96, 0.97, 0.975, 0.98, 0.985, 0.99, 0.995, 0.998, 1.0]),
#         'auc': 0.860
#     },
#     'Virchow2': {
#         'fpr': np.array([0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5,
#                          0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0]),
#         'tpr': np.array([0.0, 0.28, 0.48, 0.62, 0.71, 0.78, 0.83, 0.86, 0.89, 0.91, 0.93,
#                          0.94, 0.95, 0.96, 0.97, 0.975, 0.98, 0.985, 0.99, 0.995, 1.0]),
#         'auc': 0.838
#     },
#     'Virchow': {
#         'fpr': np.array([0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5,
#                          0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0]),
#         'tpr': np.array([0.0, 0.42, 0.62, 0.74, 0.81, 0.86, 0.89, 0.92, 0.94, 0.95, 0.96,
#                          0.97, 0.975, 0.98, 0.985, 0.99, 0.992, 0.995, 0.997, 0.999, 1.0]),
#         'auc': 0.888
#     },
#     'UNI': {
#         'fpr': np.array([0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5,
#                          0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0]),
#         'tpr': np.array([0.0, 0.22, 0.42, 0.56, 0.66, 0.74, 0.79, 0.83, 0.86, 0.89, 0.91,
#                          0.93, 0.94, 0.95, 0.96, 0.97, 0.975, 0.98, 0.985, 0.99, 1.0]),
#         'auc': 0.818
#     },
#     'ours': {
#         'fpr': np.array([0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5,
#                          0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0]),
#         'tpr': np.array([0.0, 0.38, 0.58, 0.70, 0.78, 0.84, 0.88, 0.91, 0.93, 0.94, 0.95,
#                          0.96, 0.97, 0.975, 0.98, 0.985, 0.99, 0.992, 0.995, 0.997, 1.0]),
#         'auc': 0.870
#     }
# }

fig, ax_main = plt.subplots(figsize=(8, 8), facecolor='white')
ax_main.set_facecolor('white')

# models = ['ours', 'Virchow', 'CRISP', 'Virchow2', 'UNI']

for i, model in enumerate(models):
    data = roc_data[model]
    ax_main.plot(data['tpr'], data['fpr'],
                 color=colors[i],
                 linewidth=3,
                 label=f"{model.upper()} AUC={data['auc']:.3f}",
                 alpha=0.9)

# 设置主图
ax_main.set_xlabel('1 - Specificity', fontsize=13, fontweight='normal')
ax_main.set_ylabel('Sensitivity', fontsize=13, fontweight='normal')
ax_main.legend(loc='lower right', fontsize=10, framealpha=1.0,
               edgecolor='black', fancybox=False)
ax_main.grid(False)
ax_main.set_xlim([0.0, 1.0])
ax_main.set_ylim([0.0, 1.0])
ax_main.set_aspect('equal')

# 设置刻度
ax_main.set_xticks([0.00, 0.25, 0.50, 0.75, 1.00])
ax_main.set_yticks([0.00, 0.25, 0.50, 0.75, 1.00])
ax_main.tick_params(labelsize=11)

# 设置边框
ax_main.spines['top'].set_visible(False)
ax_main.spines['right'].set_visible(False)
ax_main.spines['left'].set_edgecolor('black')
ax_main.spines['left'].set_linewidth(1.5)
ax_main.spines['bottom'].set_edgecolor('black')
ax_main.spines['bottom'].set_linewidth(1.5)

# 创建嵌入式子图（放大图）- [left, bottom, width, height]
ax_inset = fig.add_axes([0.6, 0.55, 0.3, 0.3])
ax_inset.set_facecolor('white')

# 在子图中绘制放大的ROC曲线
for i, model in enumerate(models):
    data = roc_data[model]
    ax_inset.plot(data['tpr'], data['fpr'],
                  color=colors[i],
                  linewidth=2.5,
                  alpha=0.9)

# 设置子图样式 - 放大左上角区域
ax_inset.set_xlim([0.0, 0.3])
ax_inset.set_ylim([0.7, 1.0])
ax_inset.set_xticks([0.0, 0.1, 0.2, 0.3])
ax_inset.set_yticks([0.7, 0.8, 0.9, 1.0])
ax_inset.tick_params(labelsize=9)
ax_inset.grid(False)

# 添加边框
ax_inset.spines['top'].set_visible(False)
ax_inset.spines['right'].set_visible(False)
ax_inset.spines['left'].set_edgecolor('black')
ax_inset.spines['left'].set_linewidth(1.5)
ax_inset.spines['bottom'].set_edgecolor('black')
ax_inset.spines['bottom'].set_linewidth(1.5)

# plt.tight_layout()
plt.savefig("png_auc.png")
